import torchvision.models as models
import torch.nn as nn
from util import Parameter
from ArcFace import Arcface

class EmbeddingModel(nn.Module):
    def __init__(self):
        super(EmbeddingModel,self).__init__()
        self.backbone=models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        self.backbone.fc=nn.Linear(in_features=2048, out_features=Parameter.EMBEDDING_SIZE, bias=False)
        self.arcface=Arcface(embedding_size=Parameter.EMBEDDING_SIZE,classnum=Parameter.NUM_CLASSES)
    
    def forward(self,input):
        x=self.backbone(input)
        x=self.arcface(x)
        return x
